package com.fmync.ai.fmync.config;

import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.ExecutionException;

@Configuration
public class VectorStoreConfig {

    @Value("${qdrant.host}")
    private String host;

    @Value("${qdrant.port}")
    private int port;

    @Value("${qdrant.collection-name}")
    private String collectionName;


    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {

        QdrantClient client =
                new QdrantClient(
                        QdrantGrpcClient.newBuilder(host, port, false)
                                .build());


        Boolean exists = null;
        try {
            exists = client.collectionExistsAsync(collectionName).get();
            if(!exists){
                client
                        .createCollectionAsync(
                                collectionName,
                                Collections.VectorParams.newBuilder().setDistance(Collections.Distance.Cosine).setSize(1024).build())
                        .get();
            }
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }

        QdrantEmbeddingStore qdrantEmbeddingStore = QdrantEmbeddingStore.builder()
                .host(host)
                .port(port)
                .collectionName(collectionName)
                .build();
        return qdrantEmbeddingStore;
    }

    @Bean
    public DocumentSplitter documentSplitter() {
        return DocumentSplitters.recursive(500, 0);
    }

    @Bean
    public MessageWindowChatMemory chatMemory() {
        return MessageWindowChatMemory.withMaxMessages(10);
    }
} 